import os

import torch



def data_division(mode, now_label, task_num, name):
    file_path = name + ".txt"
    save_path = "./fg1_pg1/task_{}/".format(task_num)
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    guid_index = 1
    examples = []
    with open(file_path, encoding="utf-8") as f:
        words = []
        labels = []
        with open(save_path + "{}.txt".format(mode), "w", encoding="utf-8") as f_w:
            for line in f:
                if line.startswith("-DOCSTART-") or not line.strip():
                    if words:
                        guid_index += 1
                        words = []
                        labels = []
                        f_w.write("\n")
                else:
                    line = line.strip()
                    splits = line.split(" ")

                    if splits[0].strip():
                        words.append(splits[0])
                        f_w.write(splits[0] + "\t")
                        if len(splits) > 1:
                            lab = splits[-1].replace("\n", "")
                            lab = lab.split("-")[-1]
                            # print(lab)
                            if lab in now_label:
                                labels.append(lab)
                                f_w.write(lab + "\n")
                            else:
                                labels.append("O")
                                f_w.write("O\n")
                        else:
                            # Examples could have no label for mode = "test"
                            labels.append("O")
                            f_w.write("O\n")
    print(guid_index)
    return examples


label_lists = ['CARDINAL', 'DATE', 'EVENT', 'FAC', 'GPE', 'LANGUAGE', 'LAW', 'LOC', 'MONEY', 'NORP',
               'ORDINAL', 'ORG', 'PERCENT', 'PERSON', 'PRODUCT', 'QUANTITY', 'TIME', 'WORK_OF_ART']
per_types = 1
for now_md in ["dev", "test"]:
    for i in range(len(label_lists) // per_types):
        if now_md == "test":
            labels = label_lists[:(i+1) * per_types]
        else:
            labels = label_lists[per_types * i: per_types * (i+1)]
        data_division(
            mode=now_md,
            now_label=labels,
            task_num=i,
            name=now_md)